import asyncio  
from scapy.all import ARP, Ether, srp, IP, ICMP, sr1  
import time  
from PyQt5.QtWidgets import QApplication  
from PyQt5.QtCore import QThread, pyqtSignal  
import sys  
from qasync import QEventLoop  
import logging  
import ipaddress  
from traceroute import trace_path  
from app import Canvas, MainWindow
from port_scan import scan_ports
# 定义一个异步函数来处理发送ARP请求并接收响应  
async def async_arprequest(subnet, alive_nodes, canvas):  
    arp = ARP(pdst=subnet)  
    ether = Ether(dst="ff:ff:ff:ff:ff:ff")  
    packet = ether / arp  
    result = await asyncio.get_event_loop().run_in_executor(None, lambda: srp(packet, timeout=1, verbose=False)[0])  
    for sent, received in result:  
        if received.psrc != "0.0.0.0":  
            if received.psrc not in alive_nodes: 
                task = []
                ip = received.psrc
                task.append(asyncio.create_task(scan_ports(ip, port_range=None, specific_ports=[20, 21, 22, 23, 25, 53, 67, 68, 80, 110, 123, 143, 161, 194, 220, 389, 443, 445, 465, 514, 512, 513, 543, 554, 587, 631, 8080, 8443, 3306, 5432, 6379, 9200, 11211])))
                open_list = await asyncio.gather(*task)
                alive_nodes[ip] = open_list  
                canvas.update_alive_nodes_signal.emit(alive_nodes)  

# 定义一个处理 ARP 请求的线程  
class ARPThread(QThread):  
    update_alive_nodes_signal = pyqtSignal(dict)  

    def __init__(self, subnet, alive_nodes, canvas):  
        super().__init__()  
        self.subnet = subnet  
        self.alive_nodes = alive_nodes  
        self.canvas = canvas  

    def run(self):  
        asyncio.run(async_arprequest(self.subnet, self.alive_nodes, self.canvas))  
        self.update_alive_nodes_signal.emit(self.alive_nodes)  

# 定义一个异步函数来ping主机  
async def ping_host(ip, alive_nodes, canvas):  
    packet = IP(dst=ip) / ICMP()  
    response = await asyncio.get_running_loop().run_in_executor(None, lambda: sr1(packet, timeout=1, verbose=0))  

    if response and IP in response:  
        if ip not in alive_nodes:  
            task = []
            task.append(asyncio.create_task(scan_ports(ip, port_range=None, specific_ports=[20, 21, 22, 23, 25, 53, 67, 68, 80, 110, 123, 143, 161, 194, 220, 389, 443, 445, 465, 514, 512, 513, 543, 554, 587, 631, 8080, 8443, 3306, 5432, 6379, 9200, 11211])))
            open_list = await asyncio.gather(*task)
            alive_nodes[ip] = open_list

            canvas.update_alive_nodes_signal.emit(alive_nodes)  

async def scan_multiple_subnets(subnets, alive_nodes, canvas):  
    tasks = []  
    for base_ip, subnet_start, subnet_end in subnets:  
        for i in range(subnet_start, subnet_end + 1):  
            ip = f"{base_ip}.{i}"  
            tasks.append(asyncio.create_task(ping_host(ip, alive_nodes, canvas)))  
    
    await asyncio.gather(*tasks)  

# 主函数  
async def main(app, canvas, subnet):  
    dic = {}  
    start_time = time.time()  
     
    # 这里改用线程去处理 ARP 请求  
    arp_thread = ARPThread(subnet, dic, canvas)  
    arp_thread.start()  
    
    # 创建任务列表  
    subnets_to_scan = [("61.132.163", 1, 200)]  
    tasks1 = [asyncio.create_task(scan_multiple_subnets(subnets_to_scan, dic, canvas))] 

    await asyncio.gather(*tasks1)  

    ip_list = list(dic.keys())  
    targets_list = {}  
    
    # 处理目标 IP 和 TTL 值的逻辑  
    for ip in ip_list:  
        subnet_mask = '255.255.255.0'  
        network = ipaddress.IPv4Network(f"{ip}/{subnet_mask}", strict=False)  
        if network not in targets_list:  
            targets_list[network] = []  
        targets_list[network].append(ip)  

    host_list = [(hosts[0], 20) for hosts in targets_list.values()]  

    scan_list = []  
    for target in host_list:  
        target_ip, max_hops = target  
        hop_new = await trace_path(target_ip, max_hops)  
        for ip in hop_new:  
            if ip == "*":  
                continue  
            elif ip == target_ip:  
                break  
            else:  
                more_sub = ip.split('.')  
                base_ip = '.'.join(more_sub[:-1])  
                scan_list.append((str(base_ip), 1, 255))  
                canvas.update_trace_signal.emit(scan_list) 
    await asyncio.gather(*(asyncio.create_task(scan_multiple_subnets(scan_list, dic, canvas)) for _ in scan_list))  
    
    # 发送更新 trace 的信号   
    end_time = time.time()  
    logging.info(f"Scan completed in {end_time - start_time:.2f} seconds")  

def start_main():
    asyncio.ensure_future(main(app, canvas, subnet))

# 运行主函数  
if __name__ == "__main__":  
    app = QApplication(sys.argv)  
    canvas = Canvas()  
    mainWin = MainWindow(canvas)  
    mainWin.show()  

    loop = QEventLoop(app)  
    asyncio.set_event_loop(loop)  

    # 启动异步扫描任务
    subnet = mainWin.ip_input.text()
    subnet = subnet+"/24"
    mainWin.startButton.clicked.connect(start_main)

    with loop:  
        sys.exit(loop.run_forever())